#include <graphviz.hpp>
#include <AST.hpp>
#include <sstream>
#include "llvm/Support/raw_ostream.h"

namespace graphviz {

    GraphBase &GraphBase::addAttr(std::string attr) {
        if (attr.size() == 0) return *this;
        attrList.emplace_back(move(attr));
        return *this;
    }

    GraphBase &GraphBase::addLabel(std::string label) {
        if (label.size() == 0) return *this;
        return addAttr("label=\"" + label + "\"");
    }

    void GraphBase::printAttr(std::ostream &os) const {
        os << " [ ";
        for (const auto &str : attrList) os << str << "; ";
        os << "]";
    }

    GraphNode::GraphNode(size_t idx) : id("node_" + std::to_string(idx)) {}

    GraphDiEdge::GraphDiEdge(const GraphNode &from, const GraphNode &to)
        : from(from), to(to) {}

} // namespace graphviz

namespace graphviz {
    using std::to_string;

    GraphGenerator::GraphGenerator(std::string outputFilename)
        : os(outputFilename) {
        os << "strict digraph {\n";
    }

    GraphGenerator::~GraphGenerator() {
        os << "}\n";
        os.flush();
        os.close();
    }

    GraphGenerator &GraphGenerator::get() {
        static GraphGenerator graphGenerator("graph.dot");
        return graphGenerator;
    }

    /*
    GraphGenerator &
    GraphGenerator::pushBack(const Expr *expr, std::string label) {
        if (expr == nullptr) return *this;
        visitQueue.emplace(expr, currentParentId, move(label));
        return *this;
    }

    void GraphGenerator::run() {
        while (!visitQueue.empty()) {
            auto &front = visitQueue.front();
            const Expr *expr = std::get<0>(front);
            currentParentId = move(std::get<1>(front));
            currentEdgeLabel = move(std::get<2>(front));
            visitQueue.pop();
            expr->accept(*this);
        }
    }
    */

    size_t GraphGenerator::nodeCount = 0;

    GraphNode GraphGenerator::newNode() {
        return GraphNode(nodeCount++ /*, currentParentId, currentEdgeLabel*/);
    }

    void graphviz::GraphGenerator::handleChild(
        const GraphNode &parent, const Expr *child, const std::string &hint) {
        if (child == nullptr) return;
        GraphNode childNode = child->accept(*this).value();
        GraphDiEdge edge = GraphDiEdge(parent, childNode);
        edge.addLabel(hint);
        os << edge;
    }

    VisitResult GraphGenerator::visitBreakAST(const Break *expr) {
        GraphNode node = newNode();
        node.addLabel("break");
        os << node;
        return node;
    }

    VisitResult GraphGenerator::visitContinueAST(const Continue *expr) {
        GraphNode node = newNode();
        node.addLabel("continue");
        os << node;
        return node;
    }

    VisitResult GraphGenerator::visitForAST(const For *expr) {
        GraphNode node = newNode();
        node.addLabel("for");
        os << node;
        handleChild(node, expr->init, "init");
        handleChild(node, expr->condition, "condition");
        handleChild(node, expr->step, "step");
        handleChild(node, expr->body, "body");
        return node;
    }

    VisitResult GraphGenerator::visitIfAST(const If *expr) {
        GraphNode node = newNode();
        node.addLabel("if");
        os << node;
        handleChild(node, expr->condition, "condition");
        handleChild(node, expr->thenExpr, "then");
        handleChild(node, expr->elseExpr, "else");
        return node;
    }

    VisitResult GraphGenerator::visitReturnAST(const Return *expr) {
        GraphNode node = newNode();
        node.addLabel("return");
        os << node;
        handleChild(node, expr->returnValue, "return");
        return node;
    }

    VisitResult GraphGenerator::visitWhileAST(const While *expr) {
        GraphNode node = newNode();
        node.addLabel("while");
        os << node;
        handleChild(node, expr->condition, "condition");
        handleChild(node, expr->body, "body");
        return node;
    }

    VisitResult GraphGenerator::visitCallAST(const Call *expr) {
        GraphNode node = newNode();
        node.addLabel("call\\n<ID>" + expr->funName);
        os << node;
        for (unsigned i = 0; i < expr->args.size(); ++i) {
            handleChild(node, expr->args[i], "arg" + to_string(i));
        }
        return node;
    }

    VisitResult GraphGenerator::visitFunctionAST(const FunctionAST *expr) {
        GraphNode node = newNode();
        node.addLabel("function\\ndefinition");
        os << node;
        handleChild(node, expr->proto, "prototype");
        handleChild(node, expr->body, "body");
        return node;
    }

    VisitResult GraphGenerator::visitPrototypeAST(const PrototypeAST *expr) {
        GraphNode node = newNode();
        node.addLabel("prototype\\n<ID>" + expr->funName);
        os << node;
        for (unsigned i = 0; i < expr->paraNames.size(); ++i) {
            std::string label("<");
            llvm::raw_string_ostream ss(label);
            expr->paraTypes[i]->print(ss);
            label += (expr->paraIsRefs[i] ? "> &" : "> ") + expr->paraNames[i];
            GraphNode paraNode = newNode();
            paraNode.addLabel(label);
            GraphDiEdge edge(node, paraNode);
            edge.addLabel("para" + to_string(i));
        }
        return node;
    }

    VisitResult GraphGenerator::visitBoolAST(const Bool *expr) {
        GraphNode node = newNode();
        node.addLabel("<Bool>\\n" + expr->val ? "true" : "false");
        os << node;
        return node;
    }

    VisitResult GraphGenerator::visitCharAST(const Char *expr) {
        GraphNode node = newNode();
        node.addLabel("<Char>\\n" + to_string(expr->val));
        os << node;
        return node;
    }

    VisitResult GraphGenerator::visitFloatAST(const Float *expr) {
        GraphNode node = newNode();
        node.addLabel("<Float>\\n" + to_string(expr->val));
        os << node;
        return node;
    }

    VisitResult GraphGenerator::visitIntAST(const Int *expr) {
        GraphNode node = newNode();
        node.addLabel("<Int>\\n" + to_string(expr->val));
        os << node;
        return node;
    }

    VisitResult GraphGenerator::visitStringPtrAST(const StringPtr *expr) {
        GraphNode node = newNode();
        node.addLabel("<String>\\n" + expr->val);
        os << node;
        return node;
    }

    VisitResult GraphGenerator::visitBinaryOpAST(const BinaryOp *expr) {
        GraphNode node = newNode();
        node.addLabel("<binary op>\\n" + to_string(expr->getOpType()));
        os << node;
        handleChild(node, expr->lhs, "LHS");
        handleChild(node, expr->rhs, "RHS");
        return node;
    }

    VisitResult GraphGenerator::visitDotAST(const DotAST *expr) {
        GraphNode node = newNode();
        node.addLabel("<binary op>\\n.");
        os << node;
        handleChild(node, expr->lhs, "LHS");
        handleChild(node, expr->rhs, "RHS");
        return node;
    }

    VisitResult GraphGenerator::visitArrayPtrAST(const ArrayPtr *expr) {
        GraphNode node = newNode();
        if (expr->isDefinition) node.addLabel("array def\\n<ID>" + expr->name);
        else
            node.addLabel("array ref\\n<ID>" + expr->name);
        os << node;
        return node;
    }

    VisitResult GraphGenerator::visitStructAST(const Struct *expr) {
        GraphNode node = newNode();
        node.addLabel("struct def\\n<ID>" + expr->name);
        os << node;
        for (unsigned i = 0; i < expr->varList.size(); ++i) {
            handleChild(node, expr->varList[i], "member" + to_string(i));
        }
        return node;
    }

    VisitResult GraphGenerator::visitVarAST(const Var *expr) {
        GraphNode node = newNode();
        if (expr->isDefinition) node.addLabel("var def\\n<ID>" + expr->name);
        else
            node.addLabel("var ref\\n<ID>" + expr->name);
        os << node;
        return node;
    }
} // namespace graphviz
